# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import math
import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

FLAGS = flags.FLAGS


def do_plot_spread(df, csv_root, exchange, market):
  fig, ax1 = plt.subplots(figsize=(20, 10))

  ax2 = ax1.twinx()
  ax1.tick_params(axis='x', rotation=45)
  symbols = df['symbol']
  plt.title('aggregated bid_ask_spread bp VS volume in %s 20180802-20180806' % exchange)
  plt.xticks(df.index, symbols)
  ax1.plot(df.index, df['avg_bid_ask_spread/vwap_bp_mean'], color='g', marker='o')
  ax2.plot(df.index, df['total_volume*vwap_mean'], color='b', marker='o')
  tick_space = (ax1.get_yticks()[-1] - ax1.get_yticks()[0]) / 20
  tick_space = int(math.ceil(tick_space / 5)) * 5
  ax1.yaxis.set_major_locator(ticker.MultipleLocator(tick_space))

  ax1.set_ylabel('bid_ask_spread bp', color='g')
  ax2.set_ylabel('volume USD/EUR/BTC', color='b')
  plt.savefig(os.path.join(csv_root, 'aggregated_%s_20180802-20180806.png' % exchange))


def merge_spread_volume_to_csv(dfs, exchange, market, csv_root):
  extracted_dfs = []
  for df in dfs:
    extracted_df = df.loc[(df['exchange'] == exchange)].head(100)
    extracted_df = extracted_df[[
        'exchange', 'market', 'symbol', 'avg_bid_ask_spread/vwap_bp', 'total_volume*vwap'
    ]]
    extracted_dfs.append(extracted_df)

  merged_df = extracted_dfs[0]
  start_date = 20180802
  end_date = 20180802
  for extracted_df in extracted_dfs[1:]:
    merged_df = pandas.merge(merged_df,
                             extracted_df,
                             how='outer',
                             on=['exchange', 'market', 'symbol'])
    end_date += 1
    columns = ['exchange', 'market', 'symbol']
    for date in range(start_date, end_date + 1, 1):
      columns += ['avg_bid_ask_spread/vwap_bp_%d' % date, 'total_volume*vwap_%d' % date]
    merged_df.columns = columns

  merged_df.dropna(how='any', inplace=True)
  merged_df['avg_bid_ask_spread/vwap_bp_mean'] = merged_df[[
      'avg_bid_ask_spread/vwap_bp_20180802',
      'avg_bid_ask_spread/vwap_bp_20180803',
      'avg_bid_ask_spread/vwap_bp_20180804',
      'avg_bid_ask_spread/vwap_bp_20180805',
      'avg_bid_ask_spread/vwap_bp_20180806'
  ]].mean(axis=1)
  merged_df['total_volume*vwap_mean'] = merged_df[[
      'total_volume*vwap_20180802',
      'total_volume*vwap_20180803',
      'total_volume*vwap_20180804',
      'total_volume*vwap_20180805',
      'total_volume*vwap_20180806'
  ]].mean(axis=1)

  merged_df = merged_df.sort_values('total_volume*vwap_mean', ascending=False).head(30)
  print(merged_df.to_string(index=False, float_format=lambda x: "{:12.4f}".format(x)))
  merged_df.to_csv(os.path.join(csv_root, 'spread_volume_stats_%s_%s.csv' % (exchange, market)),
                   float_format='%12.6f',
                   index=False)
  merged_df = merged_df.reset_index()
  do_plot_spread(merged_df, csv_root, exchange, market)


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  spread_dfs = []
  for date in range(20180802, 20180807, 1):
    spread_df = read_csv_into_df(csv_root, 'Kraken_%d.csv.normalized' % date)
    spread_dfs.append(spread_df)

  for exchange in ('Kraken',):
    for market in ('USD',):
      merge_spread_volume_to_csv(spread_dfs, exchange, market, csv_root)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
